use crate::Embed;
use crate::client::{AsEmbeddings, ProviderClient};
use crate::embeddings::embedding::EmbeddingModelDyn;
use crate::embeddings::{EmbeddingModel, EmbeddingsBuilder};

/// A provider client with embedding capabilities.
/// Clone is required for conversions between client types.
pub trait EmbeddingsClient: ProviderClient + Clone {
    /// The type of EmbeddingModel used by the Client
    type EmbeddingModel: EmbeddingModel;

    /// Create an embedding model with the given name.
    /// Note: default embedding dimension of 0 will be used if model is not known.
    /// If this is the case, it's better to use function `embedding_model_with_ndims`
    ///
    /// # Example
    /// ```
    /// use rig::prelude::*;
    /// use rig::providers::openai::{Client, self};
    ///
    /// // Initialize the OpenAI client
    /// let openai = Client::new("your-open-ai-api-key");
    ///
    /// let embedding_model = openai.embedding_model(openai::TEXT_EMBEDDING_3_LARGE);
    /// ```
    fn embedding_model(&self, model: &str) -> Self::EmbeddingModel;

    /// Create an embedding model with the given name and the number of dimensions in the embedding generated by the model.
    ///
    /// # Example with OpenAI
    /// ```
    /// use rig::prelude::*;
    /// use rig::providers::openai::{Client, self};
    ///
    /// // Initialize the OpenAI client
    /// let openai = Client::new("your-open-ai-api-key");
    ///
    /// let embedding_model = openai.embedding_model("model-unknown-to-rig", 3072);
    /// ```
    fn embedding_model_with_ndims(&self, model: &str, ndims: usize) -> Self::EmbeddingModel;

    /// Create an embedding builder with the given embedding model.
    ///
    /// # Example with OpenAI
    /// ```
    /// use rig::prelude::*;
    /// use rig::providers::openai::{Client, self};
    ///
    /// // Initialize the OpenAI client
    /// let openai = Client::new("your-open-ai-api-key");
    ///
    /// let embeddings = openai.embeddings(openai::TEXT_EMBEDDING_3_LARGE)
    ///     .simple_document("doc0", "Hello, world!")
    ///     .simple_document("doc1", "Goodbye, world!")
    ///     .build()
    ///     .await
    ///     .expect("Failed to embed documents");
    /// ```
    fn embeddings<D: Embed>(&self, model: &str) -> EmbeddingsBuilder<Self::EmbeddingModel, D> {
        EmbeddingsBuilder::new(self.embedding_model(model))
    }

    /// Create an embedding builder with the given name and the number of dimensions in the embedding generated by the model.
    ///
    /// # Example with OpenAI
    /// ```
    /// use rig::prelude::*;
    /// use rig::providers::openai::{Client, self};
    ///
    /// // Initialize the OpenAI client
    /// let openai = Client::new("your-open-ai-api-key");
    ///
    /// let embeddings = openai.embeddings_with_ndims(openai::TEXT_EMBEDDING_3_LARGE, 3072)
    ///     .simple_document("doc0", "Hello, world!")
    ///     .simple_document("doc1", "Goodbye, world!")
    ///     .build()
    ///     .await
    ///     .expect("Failed to embed documents");
    /// ```
    fn embeddings_with_ndims<D: Embed>(
        &self,
        model: &str,
        ndims: usize,
    ) -> EmbeddingsBuilder<Self::EmbeddingModel, D> {
        EmbeddingsBuilder::new(self.embedding_model_with_ndims(model, ndims))
    }
}

pub trait EmbeddingsClientDyn: ProviderClient {
    /// Create an embedding model with the given name.
    /// Note: default embedding dimension of 0 will be used if model is not known.
    /// If this is the case, it's better to use function `embedding_model_with_ndims`
    fn embedding_model<'a>(&self, model: &str) -> Box<dyn EmbeddingModelDyn + 'a>;

    /// Create an embedding model with the given name and the number of dimensions in the embedding generated by the model.
    fn embedding_model_with_ndims<'a>(
        &self,
        model: &str,
        ndims: usize,
    ) -> Box<dyn EmbeddingModelDyn + 'a>;
}

impl<M, T> EmbeddingsClientDyn for T
where
    T: EmbeddingsClient<EmbeddingModel = M>,
    M: EmbeddingModel + 'static,
{
    fn embedding_model<'a>(&self, model: &str) -> Box<dyn EmbeddingModelDyn + 'a> {
        Box::new(self.embedding_model(model))
    }

    fn embedding_model_with_ndims<'a>(
        &self,
        model: &str,
        ndims: usize,
    ) -> Box<dyn EmbeddingModelDyn + 'a> {
        Box::new(self.embedding_model_with_ndims(model, ndims))
    }
}

impl<T> AsEmbeddings for T
where
    T: EmbeddingsClientDyn + Clone + 'static,
{
    fn as_embeddings(&self) -> Option<Box<dyn EmbeddingsClientDyn>> {
        Some(Box::new(self.clone()))
    }
}
